import datetime
import numpy as np
import pandas as pd


#twap信号  考虑到使用优先，设计delay 放回N日前的统计值，retrun的是  当前timekey下的obj的  delay期前的值
class twap_ma:
    def __init__(self,key,timekey,freq="%H%M",peroid=20,ishismode=False,delay=0,calcmode=0,fillna=0):
        """
        Args:
            key:统计关键字
            timekey: 时间字段
            freq: 时间格式  %H:%M
            peroid: 统计周期
            ishismode: 是否历史模式
            delay: 延迟使用数据个数
            calcmode: 计算模式：对0，na这种无效数据的处理：
                0：忽略na,与0数据；最多peroid个数据，少于也计算；
                1：忽略na,与0数据 严格peroid个数据才计算，
                2：忽略na,0为正常数据，严格peroid个数据才计算，
                3.na数据，0数据做填充，参数fillna,
                4:0数据为正常数据，na数据做填充， 参数fillna
            fillna: na值处理模式  在calcmode>2时有效
        """
        self.freq=freq
        self.peroid=peroid
        self.data_his={}
        self.key=key
        self.timekey=timekey
        self.data={}
        self.delay=delay
        self.data_size=delay+2
        self.ishismode=ishismode
        self.fillna=fillna
        if calcmode in [0,1]:
            self.checkdata=self.check_data01
        elif calcmode==2:
            self.checkdata = self.check_data2
        elif calcmode==3:
            self.checkdata = self.check_data3
        else:#4
            self.checkdata = self.check_data4
        if calcmode==0:
            self.oncalc=self.oncalc0
        else:
            self.oncalc=self.oncalc1
    def check_data01(self,d,key):
        if pd.isna(d):
            return None
        elif d==0:
            return None
        else:
            return d
    def check_data2(self,d,key):
        if pd.isna(d):
            return None
        else:
            return d
    def check_data3(self,d,key):
        if pd.isna(d):
            if self.fillna == "pre":
                if self.data_his[key]:
                    d = self.data_his[key][-1]
                else:
                    d = None
            else:
                d = self.fillna
        return d
    def check_data4(self,d,key):
        if pd.isna(d):
            if self.fillna=="pre":
                if self.data_his[key]:
                    d=self.data_his[key][-1]
                else:
                    d=None
            else:
                d=self.fillna
        elif d==0:
            if self.fillna=="pre":
                if self.data_his[key]:
                    d=self.data_his[key][-1]
                else:
                    d=None
            else:
                d=self.fillna
        return d
    def oncalc0(self,b):
        rst=None
        d=b[self.key]
        tk = b[self.timekey]
        key = tk.strftime(self.freq)
        d=self.checkdata(d,key)
        if d is not None:
            if key not in self.data_his.keys():
                self.data_his[key]=[]
            self.data_his[key].append(d)
            if len(self.data_his[key])>self.peroid:
                self.data_his[key].pop(0)
            a=np.mean(self.data_his[key])
            day = datetime.datetime(tk.year, tk.month, tk.day)
            if day not in self.data.keys():
                self.data[day] = {}
            self.data[day][key] = a
            if self.delay>0:
                day=day+datetime.timedelta(days=self.delay)
            if day in self.data.keys():
                if key in self.data[day].keys():
                    rst=self.data[day][key]
        if not self.ishismode:
            days = list(self.data.keys())
            if len(days)>self.data_size:
                day=days[0]
                self.data.pop(day)
        return rst
    def oncalc1(self,b):
        rst=None
        d=b[self.key]
        tk = b[self.timekey]
        key = tk.strftime(self.freq)
        d=self.checkdata(d,key)
        if d is not None:
            if key not in self.data_his.keys():
                self.data_his[key]=[]
            self.data_his[key].append(d)
            if len(self.data_his[key])>self.peroid:
                self.data_his[key].pop(0)
            if len(self.data_his[key]) == self.peroid:
                a=np.mean(self.data_his[key])
                day = datetime.datetime(tk.year, tk.month, tk.day)
                if day not in self.data.keys():
                    self.data[day] = {}
                self.data[day][key] = a
                if self.delay>0:
                    day=day+datetime.timedelta(days=self.delay)
                if day in self.data.keys():
                    if key in self.data[day].keys():
                        rst=self.data[day][key]
        if not self.ishismode:
            days = list(self.data.keys())
            if len(days)>self.data_size:
                day=days[0]
                self.data.pop(day)
        return rst


